#!/usr/bin/python3
import cv2
import numpy as np
import pickle
import os

CP_OPEN = 0
frame = None
clfs = [None, None]
clf = None
predict_label = -1

current_color_pointer = 0
color = ['White', 'Red', 'Green', 'Yellow', 'Orange', 'Blue']
bgr_color = [(255, 255, 255), (0, 0, 255), (0, 255, 0), (0, 255, 255), (0, 127, 255), (255, 0, 0)]
color2bgr = dict(zip(color, bgr_color))

train_data = []
train_label = []

window_name = "get_color_sample"

def MouseHandler(event, x, y, flags, param):
    global frame
    global clf
    global predict_label

    global train_label
    global train_data

    # do predict
    if clf is not None:
        bgr_pix = frame[y:y+1, x:x+1]
        hsv_pix = cv2.cvtColor(bgr_pix, cv2.COLOR_BGR2HSV_FULL)
        lab_pix = cv2.cvtColor(bgr_pix, cv2.COLOR_BGR2LAB)
        whole_vec = np.concatenate((bgr_pix, hsv_pix, lab_pix), axis = 2).reshape(1, 9)

        predict_label = int(clf.predict(whole_vec))


    if event == cv2.EVENT_LBUTTONDOWN :
        bgr_pix = frame[y:y+1, x:x+1]
        hsv_pix = cv2.cvtColor(bgr_pix, cv2.COLOR_BGR2HSV_FULL)
        lab_pix = cv2.cvtColor(bgr_pix, cv2.COLOR_BGR2LAB)

        whole_vec = np.concatenate((bgr_pix, hsv_pix, lab_pix), axis = 2)
        # print(whole_vec[0][0].shape)
        train_label.append(current_color_pointer)
        train_data.append(whole_vec)
        # print(whole_vec, current_color_pointer)
    
    return 



if __name__ == "__main__":
    # 尝试加载clf
    clfs_file = None
    clfs = [None, None]
    try:
        clfs_file = open("now.model", 'rb')
        clfs = pickle.load(clfs_file)
        clfs_file.close()
        len(clfs)   # in case old format
    except Exception as e:
        print("[Warning] file 'now.model' not found or wrong format of 'now.model'.")
        clfs = [None, None]
        clfs_file = open("now.model", 'wb')
        pickle.dump(clfs, clfs_file)
        clfs_file.close()
        import time
        time.sleep(1.7)
    clf = clfs[CP_OPEN]

    try:
        datas_file = open("train_data.bin", 'rb')
        labels_file = open("train_label.bin", 'rb')
        train_datas = pickle.load(datas_file)
        train_labels = pickle.load(labels_file)
        datas_file.close()
        labels_file.close()
        if len(train_datas) != 2 or len(train_labels) != 2:
            print("[Warning] train-data files are of wrong format")
            train_datas = [[], []]
            train_labels = [[], []]
    except:
        print("[Warning] train-data files not found")
        train_datas = [[], []]
        train_labels = [[], []]
        import time
        time.sleep(1.7)

    # 记住，python中是引用赋值，指向的就是大的数组里的东西
    train_data = train_datas[CP_OPEN]
    train_label = train_labels[CP_OPEN]

    cam_idx_file = open("../../configs/cam_idx.bin", 'rb')
    cam_idx_list = pickle.load(cam_idx_file)
    #print(cam_idx_list)
    #input()
    cam_idx_file.close()

    cp = cv2.VideoCapture(cam_idx_list[CP_OPEN])
    cp.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))
    cv2.namedWindow(window_name)
    cv2.setMouseCallback(window_name, MouseHandler)
    while True:
        _, frame = cp.read()
        if not _:
            continue
        cv2.rectangle(frame, (10, 10), (100, 100), color2bgr[color[current_color_pointer]], 27)
        cv2.putText(frame, "Should be %s"%color[current_color_pointer],  (120, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, color2bgr[color[current_color_pointer]], 2)

        if predict_label >= 0:
            cv2.putText(frame, "Get %s"%color[predict_label], (10, frame.shape[0]), cv2.FONT_HERSHEY_SIMPLEX, 1.5, color2bgr[color[predict_label]], 2)

        print(train_label)
        cv2.imshow(window_name, frame)
        key = cv2.waitKey(1)
        if key == ord('s'): 
            # saving
            datas_file = open("train_data.bin", 'wb')
            labels_file = open("train_label.bin", 'wb')
            pickle.dump(train_datas, datas_file)
            pickle.dump(train_labels, labels_file)
            datas_file.close()
            labels_file.close()
            # retrain
            os.system("python3 Train.py %d"%CP_OPEN)
            cv2.putText(frame, "Detect Saving, re-training the model",  (120, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color2bgr[color[current_color_pointer]], 2)
            cv2.putText(frame, "Result is shown in the terminal",  (120, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color2bgr[color[current_color_pointer]], 2)
            # reload the clf
            clfs_file = open("now.model", 'rb')
            clfs = pickle.load(clfs_file)
            clfs_file.close()
            clf = clfs[CP_OPEN]

            cv2.putText(frame, "Classifier reload successfully",  (120, 200), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color2bgr[color[current_color_pointer]], 2)
            cv2.putText(frame, "Press any key to continue...",  (120, 250), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color2bgr[color[current_color_pointer]], 2)
            cv2.imshow(window_name, frame)
            cv2.waitKey()
        elif key == ord('p'):
            current_color_pointer += 1
            current_color_pointer %= 6
        elif key == ord('d'):
            if len(train_data) > 0:
                del train_data[-1]
                del train_label[-1]
        elif key == ord('q'):
            break
        elif key == ord('c'):
            CP_OPEN += 1
            CP_OPEN %= 2
            print(cam_idx_list[CP_OPEN], CP_OPEN)
            cp.open(cam_idx_list[CP_OPEN])

            clf = clfs[CP_OPEN]
            train_data = train_datas[CP_OPEN]
            train_label = train_labels[CP_OPEN]
